import scipy.io as sio
import numpy as np

def getData():
    data=sio.loadmat('shreejoy1.mat') # Spike times for 43 cells, models and data.  
    # Models contain 200 trials each; data contains 40 trials each.  
    return data
        
"""We want these initialization functions to be methods of a class for data/models."""        
        
def initModel(data): # Takes a proprietary data structure (Shreejoy's) and returns the model data.  
    d=data['simSpikesAll'][0] # Cells x 1 x Max Spikes x Trials
    # Max Spikes is the maximum number observed in any trial for that cell.  
    model=dict(moe='model',cells={})
    nCells=len(d)
    for cell in range(nCells):       
        nTrials=len(d[cell][0][0])
        model['cells'][cell]=dict(cellID='Some Identifier',trials={})
        for trial in range(nTrials):
            spikeTimes=d[cell][0][:,trial] # Spike times on this trial, including zero padding.  
            model['cells'][cell]['trials'][trial]=np.array(spikeTimes[np.where(spikeTimes)]) # Spike times on this trial, without zero padding (only the times)
    return model
    
def initExperiment(data): # Takes a proprietary data structure (Shreejoy's) and returns the experiment data.  
    d=data['testSpikesAll'][0] # Cells x 1 x Max Spikes x Trials
    # Max Spikes is the maximum number observed in any trial for that cell.  
    experiment=dict(moe='experiment',cells={})
    nCells=len(d)
    for cell in range(nCells):       
        nTrials=len(d[cell][0][0])
        experiment['cells'][cell]=dict(cellID='Some Identifier',trials={})
        for trial in range(nTrials):
            spikeTimes=d[cell][0][:,trial] # Spike times on this trial, including zero padding.  
            experiment['cells'][cell]['trials'][trial]=np.array(spikeTimes[np.where(spikeTimes)]) # Spike times on this trial, without zero padding (only the times)
    return experiment

class cell: 
    def __init__(self,moe,index=0): # Needs a model or experiment (moe).  
        self.moe=moe
        self.id=moe['cells'][index]['cellID']
        self.trials=moe['cells'][index]['trials']
    def firingRate(self):
        nTrials=len(self.trials)
        nSpikes=0
        for trial in range(nTrials):
            nSpikes+=len(self.trials[trial])
        return float(nSpikes)/float(nTrials)
        
class HasFiringRate(object):
    def firingRate(self):
        """Returns the firing rate of the cell, in Hz."""
        pass
    
class HasFiringRateVariance(object):
    def firingVariance(self):
        """Returns the firing rate variance of the cell, in Hz^2."""
        pass
    
class GenSpikeTimes(object):
    def spikeTrain(self):
        """Returns a structure containing the spike train."""        
        pass

class MyCell(HasFiringRate, HasFiringRateVariance):
    def firingRate(self):
        return 10.0
        
